#pragma once
#include <highgui.h>
#include "cv.h"


namespace Calibration
{
	
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

	double compute_reprojection_error( const CvMat* object_points,
		const CvMat* rot_vects, const CvMat* trans_vects,
		const CvMat* camera_matrix, const CvMat* dist_coeffs,
		const CvMat* image_points, const CvMat* point_counts,
		CvMat* per_view_errors )
	{
		CvMat* image_points2 = cvCreateMat( image_points->rows,
			image_points->cols, image_points->type );
		int i, image_count = rot_vects->rows, points_so_far = 0;
		double total_err = 0, err;

		for( i = 0; i < image_count; i++ )
		{
			CvMat object_points_i, image_points_i, image_points2_i;
			int point_count = point_counts->data.i[i];
			CvMat rot_vect, trans_vect;

			cvGetCols( object_points, &object_points_i,
				points_so_far, points_so_far + point_count );
			cvGetCols( image_points, &image_points_i,
				points_so_far, points_so_far + point_count );
			cvGetCols( image_points2, &image_points2_i,
				points_so_far, points_so_far + point_count );
			points_so_far += point_count;

			cvGetRow( rot_vects, &rot_vect, i );
			cvGetRow( trans_vects, &trans_vect, i );

			cvProjectPoints2( &object_points_i, &rot_vect, &trans_vect,
				camera_matrix, dist_coeffs, &image_points2_i,
				0, 0, 0, 0, 0 );
			err = cvNorm( &image_points_i, &image_points2_i, CV_L1 );
			if( per_view_errors )
				per_view_errors->data.db[i] = err/point_count;
			total_err += err;
		}

		cvReleaseMat( &image_points2 );
		return total_err/points_so_far;
	}


	int run_calibration( CvSeq* image_points_seq, CvSize img_size, CvSize board_size,
		float square_size, float aspect_ratio, int flags,
		CvMat* camera_matrix, CvMat* dist_coeffs, CvMat** extr_params,
		CvMat** reproj_errs, double* avg_reproj_err )
	{
		int code;
		int image_count = image_points_seq->total;
		int point_count = board_size.width*board_size.height;
		CvMat* image_points = cvCreateMat( 1, image_count*point_count, CV_32FC2 );
		CvMat* object_points = cvCreateMat( 1, image_count*point_count, CV_32FC3 );
		CvMat* point_counts = cvCreateMat( 1, image_count, CV_32SC1 );
		CvMat rot_vects, trans_vects;
		int i, j, k;
		CvSeqReader reader;
		cvStartReadSeq( image_points_seq, &reader );

		// initialize arrays of points
		for( i = 0; i < image_count; i++ )
		{
			CvPoint2D32f* src_img_pt = (CvPoint2D32f*)reader.ptr;
			CvPoint2D32f* dst_img_pt = ((CvPoint2D32f*)image_points->data.fl) + i*point_count;
			CvPoint3D32f* obj_pt = ((CvPoint3D32f*)object_points->data.fl) + i*point_count;

			for( j = 0; j < board_size.height; j++ )
				for( k = 0; k < board_size.width; k++ )
				{
					*obj_pt++ = cvPoint3D32f(j*square_size, k*square_size, 0);
					*dst_img_pt++ = *src_img_pt++;
				}
				CV_NEXT_SEQ_ELEM( image_points_seq->elem_size, reader );
		}

		cvSet( point_counts, cvScalar(point_count) );

		*extr_params = cvCreateMat( image_count, 6, CV_32FC1 );
		cvGetCols( *extr_params, &rot_vects, 0, 3 );
		cvGetCols( *extr_params, &trans_vects, 3, 6 );

		cvZero( camera_matrix );
		cvZero( dist_coeffs );

		if( flags & CV_CALIB_FIX_ASPECT_RATIO )
		{
			camera_matrix->data.db[0] = aspect_ratio;
			camera_matrix->data.db[4] = 1.;
		}

		cvCalibrateCamera2( object_points, image_points, point_counts,
			img_size, camera_matrix, dist_coeffs,
			&rot_vects, &trans_vects, flags );

		code = cvCheckArr( camera_matrix, CV_CHECK_QUIET ) &&
			cvCheckArr( dist_coeffs, CV_CHECK_QUIET ) &&
			cvCheckArr( *extr_params, CV_CHECK_QUIET );

		*reproj_errs = cvCreateMat( 1, image_count, CV_64FC1 );
		*avg_reproj_err =
			compute_reprojection_error( object_points, &rot_vects, &trans_vects,
			camera_matrix, dist_coeffs, image_points, point_counts, *reproj_errs );

		cvReleaseMat( &object_points );
		cvReleaseMat( &image_points );
		cvReleaseMat( &point_counts );

		return code;
	}


	void save_camera_params( const char* out_filename, int image_count, CvSize img_size,
		CvSize board_size, float square_size,
		float aspect_ratio, int flags,
		const CvMat* camera_matrix, CvMat* dist_coeffs,
		const CvMat* extr_params, const CvSeq* image_points_seq,
		const CvMat* reproj_errs, double avg_reproj_err )
	{
		CvFileStorage* fs = cvOpenFileStorage( out_filename, 0, CV_STORAGE_WRITE );

		cvWriteInt( fs, "image_count", image_count );
		cvWriteInt( fs, "image_width", img_size.width );
		cvWriteInt( fs, "image_height", img_size.height );
		cvWriteInt( fs, "board_width", board_size.width );
		cvWriteInt( fs, "board_height", board_size.height );
		cvWriteReal( fs, "square_size", square_size );

		if( flags & CV_CALIB_FIX_ASPECT_RATIO )
			cvWriteReal( fs, "aspect_ratio", aspect_ratio );

		char buf[1024];
		if( flags != 0 )
		{
			sprintf( buf, "flags: %s%s%s%s",
				flags & CV_CALIB_USE_INTRINSIC_GUESS ? "+use_intrinsic_guess" : "",
				flags & CV_CALIB_FIX_ASPECT_RATIO ? "+fix_aspect_ratio" : "",
				flags & CV_CALIB_FIX_PRINCIPAL_POINT ? "+fix_principal_point" : "",
				flags & CV_CALIB_ZERO_TANGENT_DIST ? "+zero_tangent_dist" : "" );
			cvWriteComment( fs, buf, 0 );
		}

		cvWriteInt( fs, "flags", flags );

		cvWrite( fs, "camera_matrix", camera_matrix );
		cvWrite( fs, "distortion_coefficients", dist_coeffs );

		cvWriteReal( fs, "avg_reprojection_error", avg_reproj_err );
		if( reproj_errs )
			cvWrite( fs, "per_view_reprojection_errors", reproj_errs );

		if( extr_params )
		{
			cvWriteComment( fs, "a set of 6-tuples (rotation vector + translation vector) for each view", 0 );
			cvWrite( fs, "extrinsic_parameters", extr_params );
		}

		if( image_points_seq )
		{
			cvWriteComment( fs, "the array of board corners projections used for calibration", 0 );
			assert( image_points_seq->total == image_count );
			CvMat* image_points = cvCreateMat( 1, image_count*board_size.width*board_size.height, CV_32FC2 );
			cvCvtSeqToArray( image_points_seq, image_points->data.fl );

			cvWrite( fs, "image_points", image_points );
			cvReleaseMat( &image_points );
		}

		cvReleaseFileStorage( &fs );
	}


////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

class BoardCalibrator
	{
	protected:
		// common constants
		CvSize _boardSize;
		CvSize _imagesSize;
		float _squareSize;
		float _aspectRatio;
		int _flags;

		CvMemStorage* _storage;
		CvSeq* _imagePointsSeq;
		CvPoint2D32f* _imagePointsBuf;
		double _cameraData[9], _distCoeffsData[4];

		CvMat _camera;
		CvMat _distCoeffs;
		CvMat* _extrParams;
		CvMat*_reprojPerrs;

		double _avgReprojError;

	public: 
		BoardCalibrator(CvSize imagesSize)
		{
			_imagesSize = imagesSize;
			_boardSize.width = 8;
			_boardSize.height = 6;

			_flags = 0;
			_squareSize  = 1.f;
			_aspectRatio = 1.f;

			_imagePointsSeq = NULL;
			_storage = NULL;
			_imagePointsBuf = NULL;

			_camera = cvMat( 3, 3, CV_64F, _cameraData );
			_distCoeffs = cvMat( 1, 4, CV_64F, _distCoeffsData );
			
			_extrParams = NULL;
			_reprojPerrs = NULL;

			int elementSize = _boardSize.width*_boardSize.height*sizeof(_imagePointsBuf[0]);
			_storage = cvCreateMemStorage( MAX( elementSize*4, 1 << 16 ));
			
			_imagePointsBuf = (CvPoint2D32f*)cvAlloc( elementSize );
			_imagePointsSeq = cvCreateSeq( 0, sizeof(CvSeq), elementSize, _storage );
			
			_avgReprojError = 0;
		}

		void AddImage(IplImage* img)
		{			
			int count = 0;
			int found = cvFindChessboardCorners( img, _boardSize, _imagePointsBuf, &count, CV_CALIB_CB_ADAPTIVE_THRESH );

			// subpixel refinement
			IplImage* gray = cvCreateImage( _imagesSize, 8, 1 );
			cvCvtColor( img, gray, CV_BGR2GRAY );
			cvFindCornerSubPix( gray, _imagePointsBuf, count, cvSize(11,11), cvSize(-1,-1), cvTermCriteria( CV_TERMCRIT_EPS+CV_TERMCRIT_ITER, 30, 0.1 ));
			cvReleaseImage( &gray );
			
            cvSeqPush( _imagePointsSeq, _imagePointsBuf );

			//IplImage* t = cvCloneImage( img );
			//cvUndistort2( t, img , &_camera, &_distCoeffs );
			//cvReleaseImage( &t );
			cvDrawChessboardCorners( img, _boardSize, _imagePointsBuf, count, found );
		}

		void Calibrate()
		{
			int code = run_calibration( _imagePointsSeq, _imagesSize, _boardSize,
				_squareSize, _aspectRatio, _flags, &_camera, &_distCoeffs, &_extrParams,
                &_reprojPerrs, &_avgReprojError );
		}

		void Save(char* fileName)
		{
			save_camera_params( fileName, _imagePointsSeq->total, _imagesSize,
				_boardSize, _squareSize, _aspectRatio, _flags,
				&_camera, &_distCoeffs, _extrParams, _imagePointsSeq, _reprojPerrs, _avgReprojError );
		}

		virtual ~BoardCalibrator()
		{
			if( _storage )
				cvReleaseMemStorage( &_storage );
		}

	};
}